{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# IERG 6130 Assignment 2: Value Function Approximation in RL\n",
    "\n",
    "*2019-2020 2nd term, IERG 6130: Reinforcement Learning and Beyond. Department of Information Engineering, The Chinese University of Hong Kong. Course Instructor: Professor ZHOU Bolei. Assignment author: PENG Zhenghao.*\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| Student Name | Student ID |\n",
    "| :----: | :----: |\n",
    "| ZhengboWang | None |\n",
    "\n",
    "------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Welcome to the assignment 2 of our RL course. \n",
    "\n",
    "You need to go through this self-contained notebook, which contains many TODOs in part of the cells and has special `[TODO]` signs. You need to finish all TODOs. Some of them may be easy such as uncommenting a line, some of them may be difficult such as implementing a function. You can find them by searching the `[TODO]` symbol. However, we suggest you to go through the notebook step by step, which would give you a better understanding of the content.\n",
    "\n",
    "You are encouraged to add more code on extra cells at the end of the each section to investigate the problems you think interesting. At the end of the file, we left a place for you to optionally write comments (Yes, please give us rewards so we can keep improving the assignment!). \n",
    "\n",
    "Please report any code bugs to us via cuhkrlcourse@googlegroups.com or via github issue.\n",
    "\n",
    "Before you get start, remember to follow the instruction at https://github.com/cuhkrlcourse/ierg6130 to setup your environment.\n",
    "\n",
    "We will cover the following knowledge in this assignment:\n",
    "\n",
    "1. The n-step Sarsa and Q learning algorithm\n",
    "2. Linear function as the approximation of value function\n",
    "3. Feature construction\n",
    "4. Neural network based function approximation\n",
    "5. Naive implementation of Deep Q Network using pytorch\n",
    "\n",
    "As shown in the following figure, in the first section of notebook, we build a basic RL pipeline. In the second section, we implement the linear function as approximator and also introduces feature construction technique. In the third section, we implement a simple neural network simply using Numpy package. \n",
    "\n",
    "![](overview.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------\n",
    "\n",
    "## Section 1: Basic Reinforcement Learning Pipeline\n",
    "\n",
    "(5 / 100 points)\n",
    "\n",
    "In this section, we will prepare several functions for evaulation and training of RL algorithms. We will also build an `AbstractTrainer` class used as a general framework which left blanks for different function approximation methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "import gym\n",
    "import numpy as np\n",
    "import torch\n",
    "from utils import *\n",
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "def evaluate(policy, num_episodes=1, seed=0, env_name='FrozenLake8x8-v0',\n",
    "             render=False):\n",
    "    \"\"\"This function evaluate the given policy and return the mean episode \n",
    "    reward.\n",
    "    :param policy: a function whose input is the observation\n",
    "    :param num_episodes: number of episodes you wish to run\n",
    "    :param seed: the random seed\n",
    "    :param env_name: the name of the environment\n",
    "    :param render: a boolean flag indicating whether to render policy\n",
    "    :return: the averaged episode reward of the given policy.\n",
    "    \"\"\"\n",
    "    env = gym.make(env_name)\n",
    "    env.seed(seed)\n",
    "    rewards = []\n",
    "    if render: num_episodes = 1\n",
    "    for i in range(num_episodes):\n",
    "        obs = env.reset()\n",
    "        act = policy(obs)\n",
    "        ep_reward = 0\n",
    "        while True:\n",
    "            obs, reward, done, info = env.step(act)\n",
    "            act = policy(obs)\n",
    "            ep_reward += reward\n",
    "            if render:\n",
    "                env.render()\n",
    "                wait(sleep=0.05)\n",
    "            if done:\n",
    "                break\n",
    "        rewards.append(ep_reward)\n",
    "    if render:\n",
    "        env.close()\n",
    "    return np.mean(rewards)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "def run(trainer_cls, config=None, reward_threshold=None):\n",
    "    \"\"\"Run the trainer and report progress, agnostic to the class of trainer\n",
    "    :param trainer_cls: A trainer class \n",
    "    :param config: A dict\n",
    "    :param reward_threshold: the reward threshold to break the training\n",
    "    :return: The trained trainer and a dataframe containing learning progress\n",
    "    \"\"\"\n",
    "    assert inspect.isclass(trainer_cls)\n",
    "    if config is None:\n",
    "        config = {}\n",
    "    trainer = trainer_cls(config)\n",
    "    config = trainer.config\n",
    "    start = now = time.time()\n",
    "    stats = []\n",
    "    for i in range(config['max_iteration'] + 1):\n",
    "        stat = trainer.train()\n",
    "        stats.append(stat or {})\n",
    "        if i % config['evaluate_interval'] == 0 or \\\n",
    "                i == config[\"max_iteration\"]:\n",
    "            reward = trainer.evaluate(config.get(\"evaluate_num_episodes\", 50))\n",
    "            print(\"({:.1f}s,+{:.1f}s)\\tIteration {}, current mean episode \"\n",
    "                  \"reward is {}. {}\".format(\n",
    "                time.time() - start, time.time() - now, i, reward,\n",
    "                {k: round(np.mean(v), 4) for k, v in\n",
    "                 stat.items()} if stat else \"\"))\n",
    "            now = time.time()\n",
    "        if reward_threshold is not None and reward > reward_threshold:\n",
    "            print(\"In {} iteration, current mean episode reward {:.3f} is \"\n",
    "                  \"greater than reward threshold {}. Congratulation! Now we \"\n",
    "                  \"exit the training process.\".format(\n",
    "                i, reward, reward_threshold))\n",
    "            break\n",
    "    return trainer, stats\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve TODOs and remove \"pass\"\n",
    "\n",
    "default_config = dict(\n",
    "    env_name=\"CartPole-v0\",\n",
    "    max_iteration=1000,\n",
    "    max_episode_length=1000,\n",
    "    evaluate_interval=100,\n",
    "    gamma=0.99,\n",
    "    eps=0.3,\n",
    "    seed=0\n",
    ")\n",
    "\n",
    "\n",
    "class AbstractTrainer:\n",
    "    \"\"\"This is the abstract class for value-based RL trainer. We will inherent\n",
    "    the specify algorithm's trainer from this abstract class, so that we can\n",
    "    reuse the codes.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        self.config = merge_config(config, default_config)\n",
    "\n",
    "        # Create the environment\n",
    "        self.env_name = self.config['env_name']\n",
    "        self.env = gym.make(self.env_name)\n",
    "        if self.env_name == \"Pong-ram-v0\":\n",
    "            self.env = wrap_deepmind_ram(self.env)\n",
    "\n",
    "        # Apply the random seed\n",
    "        self.seed = self.config[\"seed\"]\n",
    "        np.random.seed(self.seed)\n",
    "        self.env.seed(self.seed)\n",
    "\n",
    "        # We set self.obs_dim to the number of possible observation\n",
    "        # if observation space is discrete, otherwise the number\n",
    "        # of observation's dimensions. The same to self.act_dim.\n",
    "        if isinstance(self.env.observation_space, gym.spaces.box.Box):\n",
    "            assert len(self.env.observation_space.shape) == 1\n",
    "            self.obs_dim = self.env.observation_space.shape[0]\n",
    "            self.discrete_obs = False\n",
    "        elif isinstance(self.env.observation_space,\n",
    "                        gym.spaces.discrete.Discrete):\n",
    "            self.obs_dim = self.env.observation_space.n\n",
    "            self.discrete_obs = True\n",
    "        else:\n",
    "            raise ValueError(\"Wrong observation space!\")\n",
    "\n",
    "        if isinstance(self.env.action_space, gym.spaces.box.Box):\n",
    "            assert len(self.env.action_space.shape) == 1\n",
    "            self.act_dim = self.env.action_space.shape[0]\n",
    "        elif isinstance(self.env.action_space, gym.spaces.discrete.Discrete):\n",
    "            self.act_dim = self.env.action_space.n\n",
    "        else:\n",
    "            raise ValueError(\"Wrong action space!\")\n",
    "\n",
    "        self.eps = self.config['eps']\n",
    "\n",
    "        # You need to setup the parameter for your function approximator.\n",
    "        self.initialize_parameters()\n",
    "\n",
    "    def initialize_parameters(self):\n",
    "        self.parameters = None\n",
    "        raise NotImplementedError(\n",
    "            \"You need to override the \"\n",
    "            \"Trainer._initialize_parameters() function.\")\n",
    "\n",
    "    def process_state(self, state):\n",
    "        \"\"\"Preprocess the state (observation) if necessary\"\"\"\n",
    "        processed_state = state\n",
    "        return processed_state\n",
    "\n",
    "    def compute_values(self, processed_state):\n",
    "        \"\"\"Approximate the state value of given state.\n",
    "        This is a private function.\n",
    "        Note that you should NOT preprocess the state here.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError(\"You need to override the \"\n",
    "                                  \"Trainer.compute_values() function.\")\n",
    "\n",
    "    def compute_action(self, processed_state, eps=None):\n",
    "        \"\"\"Compute the action given the state. Note that the input\n",
    "        is the processed state.\"\"\"\n",
    "\n",
    "        values = self.compute_values(processed_state)\n",
    "        assert values.ndim == 1, values.shape\n",
    "\n",
    "        if eps is None:\n",
    "            eps = self.eps\n",
    "\n",
    "        # [TODO] Implement the epsilon-greedy policy here. We have `eps`\n",
    "        #  probability to choose a uniformly random action in action_space,\n",
    "        #  otherwise choose action that maximizes the values.\n",
    "        # Hint: Use the function of self.env.action_space to sample random\n",
    "        # action.\n",
    "        action = None\n",
    "        if np.random.rand(1) > eps:\n",
    "            action = np.argmax(values)\n",
    "        else:\n",
    "            action = np.random.choice(self.act_dim)\n",
    "        return action\n",
    "\n",
    "    def evaluate(self, num_episodes=50, *args, **kwargs):\n",
    "        \"\"\"Use the function you write to evaluate current policy.\n",
    "        Return the mean episode reward of 50 episodes.\"\"\"\n",
    "        policy = lambda raw_state: self.compute_action(\n",
    "            self.process_state(raw_state), eps=0.0)\n",
    "        result = evaluate(policy, num_episodes, seed=self.seed,\n",
    "                          env_name=self.env_name, *args, **kwargs)\n",
    "        return result\n",
    "\n",
    "    def compute_gradient(self, *args, **kwargs):\n",
    "        \"\"\"Compute the gradient.\"\"\"\n",
    "        raise NotImplementedError(\n",
    "            \"You need to override the Trainer.compute_gradient() function.\")\n",
    "\n",
    "    def apply_gradient(self, *args, **kwargs):\n",
    "        \"\"\"Compute the gradient\"\"\"\n",
    "        raise NotImplementedError(\n",
    "            \"You need to override the Trainer.apply_gradient() function.\")\n",
    "\n",
    "    def train(self):\n",
    "        \"\"\"Conduct one iteration of learning.\"\"\"\n",
    "        raise NotImplementedError(\"You need to override the \"\n",
    "                                  \"Trainer.train() function.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average episode reward for a random policy in 500 episodes in CartPole-v0:  22.068\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "class TestTrainer(AbstractTrainer):\n",
    "    \"\"\"This class is used for testing. We don't really train anything.\"\"\"\n",
    "    def compute_values(self, state):\n",
    "        return np.random.random_sample(size=self.act_dim)\n",
    "    def initialize_parameters(self):\n",
    "        self.parameters = np.random.random_sample(size=(self.obs_dim, self.act_dim))\n",
    "    \n",
    "t = TestTrainer(dict(env_name=\"CartPole-v0\"))\n",
    "obs = t.env.observation_space.sample()\n",
    "processed = t.process_state(obs)\n",
    "assert processed.shape == (4, )\n",
    "assert np.all(processed == obs)\n",
    "# Test compute_action\n",
    "values = t.compute_values(processed)\n",
    "correct_act = np.argmax(values)\n",
    "assert t.compute_action(processed, eps=0) == correct_act\n",
    "print(\"Average episode reward for a random policy in 500 episodes in CartPole-v0: \",\n",
    "      t.evaluate(num_episodes=500))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 2: Linear function approximation\n",
    "\n",
    "In this section, we implement a simple linear function whose input is the state (or the processed state) and output is the state-action values.\n",
    "\n",
    "First, we implement a `LinearTrainer` class which implements (1). Linear function approximation and (2). n-step semi-gradient method to update the linear function.\n",
    "\n",
    "Then we further implement a `LinearTrainerWithFeatureConstruction` class which processes the input state and provide polynomial features which increase the utility of linear function approximation.\n",
    "\n",
    "We refer the Chapter 9.4 (linear method), 9.5 (feature construction), and 10.2 (n-step semi-gradient method) of the RL textbook to you.\n",
    "\n",
    "In this section, we leverage the n-step semi-gradient Sarsa as the training algorithm. What is the \"correct value\" of an action $a_t$ at state $s_t$ in one-step case? We consider it is $r_t + \\gamma Q(s_{t+1}, a_{t+1})$ and thus lead to the TD error $TD = r_t + \\gamma Q(s_{t+1}, a_{t+1}) - Q(s_t, a_t)$. In n-step case, the target value of Q is natually extended:\n",
    "\n",
    "$$Q(s_t, a_t) = \\sum_{i=t}^{t+n-1}\\gamma^{i-t}r_i + \\gamma^n Q(s_{t+n}, a_{t+n})$$\n",
    "\n",
    "We follow the pipeline depicted in Chapter 10.2 (page 247) of the textbook to implement this logic. Note that notation of the time step of reward is different in this assignment and in the textbook. In textbook, the reward $R_{t+1}$ is the reward when apply action $a_{t}$ to the environment at state $s_t$. In the equation above the $r_t$ has exactly the same meaning. In the code below, we store the states, actions and rewards to a list during training. You need to make sure the indices of these list, like the `tau` in  `actions[tau]` has the correct meaning.\n",
    "\n",
    "After computing the target Q value, we need to derive the gradient to update the parameters. Consider a loss function, the Mean Square Error between the target Q value and the output Q value: \n",
    "\n",
    "$$\\text{loss} = \\cfrac{1}{2}[\\sum_{i=t}^{t+n-1}\\gamma^{i-t}r_i + \\gamma^n Q(s_{t+n}, a_{t+n}) - Q(s_t, a_t)]^2$$\n",
    "\n",
    "Compute the gradient of Loss with respect to the Q function:\n",
    "\n",
    "$$\\cfrac{d \\text{loss}}{d Q} = -(\\sum_{i=t}^{t+n-1}\\gamma^{i-t}r_i + \\gamma^n Q(s_{t+n}, a_{t+n}) - Q(s_t, a_t))$$\n",
    "\n",
    "According to the chain rule, the gradient of the loss w.r.t. the parameter ($W$) is:\n",
    "\n",
    "$$\\cfrac{d \\text{loss}}{d W} = -(\\sum_{i=t}^{t+n-1}\\gamma^{i-t}r_i + \\gamma^n Q(s_{t+n}, a_{t+n}) - Q(s_t, a_t))\\cfrac{d Q}{d W}$$\n",
    "\n",
    "To minimize the loss, we only need to descent the gradient:\n",
    "\n",
    "$$W = W - lr \\cfrac{d \\text{loss}}{d W}$$\n",
    "\n",
    "wherein $lr$ is the learning rate. Therefore, in conclusion the update rule of parameters is:\n",
    "\n",
    "$$W = W + lr (\\sum_{i=t}^{t+n-1}\\gamma^{i-t}r_i + \\gamma^n Q(s_{t+n}, a_{t+n}) - Q(s_t, a_t))\\cfrac{d Q}{d W}$$\n",
    "\n",
    "In the following codes, we denote $G = \\sum_{i=t}^{t+n-1}\\gamma^{i-t}r_i + \\gamma^n Q(s_{t+n}, a_{t+n})$ and will compute $dQ / dW$ according to the form of the approximator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 2.1: Basics\n",
    "\n",
    "(30 / 100 points)\n",
    "\n",
    "We want to approximate the state-action values. That is, the expected return when applying action $a_t$ in state $s_t$. Linear function approximates state-action value function by the inner product between a parameter matrix $W$ and the input state vector $s$:\n",
    "\n",
    "$$v(s, W) = W^T s$$\n",
    "\n",
    "Note that $W\\in \\mathbb R^{(O, A)}$ and $s \\in \\mathbb R^{(O, 1)}$, wherein O is the observation (state) dimensions, namely the `self.obs_dim` and A is the action dimension, namely the `self.act_dim`. Each action corresponding to one state-action values $Q(s, a)$.\n",
    "\n",
    "Note that you should finish this section **purely by Numpy without calling any other packages**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "# Build the algorithm-specify config.\n",
    "linear_approximator_config = merge_config(dict(\n",
    "    parameter_std=0.01,\n",
    "    learning_rate=0.01,\n",
    "    n=3,\n",
    "), default_config)\n",
    "\n",
    "\n",
    "class LinearTrainer(AbstractTrainer):\n",
    "    def __init__(self, config):\n",
    "        config = merge_config(config, linear_approximator_config)\n",
    "\n",
    "        # Initialize the abstract class.\n",
    "        super().__init__(config)\n",
    "\n",
    "        self.max_episode_length = self.config[\"max_episode_length\"]\n",
    "        self.learning_rate = self.config[\"learning_rate\"]\n",
    "        self.gamma = self.config[\"gamma\"]\n",
    "        self.n = self.config[\"n\"]\n",
    "\n",
    "    def initialize_parameters(self):\n",
    "        # [TODO] Initialize self.parameters, which is a two dimensional matrix,\n",
    "        #  and subjects to a normal distribution with scale\n",
    "        #  config[\"parameter_std\"].\n",
    "        std = self.config[\"parameter_std\"]\n",
    "        self.parameters = None\n",
    "        self.parameters = std * np.random.randn(self.obs_dim, self.act_dim)\n",
    "        pass\n",
    "        \n",
    "        print(\"Initialize parameters with shape: {}.\".format(\n",
    "            self.parameters.shape))\n",
    "\n",
    "    def compute_values(self, processed_state):\n",
    "        assert processed_state.ndim == 1, processed_state.shape\n",
    "        \n",
    "        # [TODO] Compute the value for each potential action. Note that you\n",
    "        #  should NOT preprocess the state here.\"\"\"\n",
    "        values = np.dot(self.parameters.T,processed_state)\n",
    "        return values\n",
    "\n",
    "    def train(self):\n",
    "        \"\"\"\n",
    "        Please implement the n-step Sarsa algorithm presented in Chapter 10.2\n",
    "        of the textbook. You algorithm should reduce the convention one-step\n",
    "        Sarsa when n = 1. That is:\n",
    "            TD = r_t + gamma * Q(s_t+1, a_t+1) - Q(s_t, a_t)\n",
    "            Q(s_t, a_t) = Q(s_t, a_t) + learning_rate * TD\n",
    "        \"\"\"\n",
    "        s = self.env.reset()\n",
    "        processed_s = self.process_state(s)\n",
    "        processed_states = [processed_s]\n",
    "        rewards = [0.0]\n",
    "        actions = [self.compute_action(processed_s)]\n",
    "        T = float(\"inf\")\n",
    "        for t in range(self.max_episode_length):\n",
    "            if t < T:\n",
    "                # [TODO]  When the termination is not reach, apply action,\n",
    "                #  process state, record state / reward / action to the\n",
    "                #  lists defined above, and deal with termination.\n",
    "                \n",
    "                act = actions[t]\n",
    "                obs,reward,done,_ = self.env.step(act)\n",
    "                processed_s = self.process_state(obs)\n",
    "                processed_states.append(processed_s)\n",
    "                rewards.append(reward)\n",
    "                if done:\n",
    "                    T = t+1\n",
    "                else:\n",
    "                    actions.append(self.compute_action(processed_s))\n",
    "\n",
    "            tau = t - self.n + 1\n",
    "            if tau >= 0:\n",
    "                gradient = self.compute_gradient(\n",
    "                    processed_states, actions, rewards, tau, T\n",
    "                )\n",
    "                self.apply_gradient(gradient)\n",
    "            if tau == T - 1:\n",
    "                break\n",
    "\n",
    "    def compute_gradient(self, processed_states, actions, rewards, tau, T):\n",
    "        \"\"\"Compute the gradient\"\"\"\n",
    "        n = self.n\n",
    "        # [TODO] Compute the approximation goal, the truth state action value\n",
    "        #  G. It is a n-step discounted sum of rewards. Refer to Chapter 10.2\n",
    "        #  of the textbook.\n",
    "        G = None\n",
    "        G = 0\n",
    "        for i in range(tau+1,min(tau+n+1,T+1)):\n",
    "            G = G + (self.gamma**(i-tau-1)) * rewards[i]\n",
    "        if tau + n < T:\n",
    "            G = G + (self.gamma**n)*self.compute_values(processed_states[tau+n])[actions[tau+n]]\n",
    "            \n",
    "        \n",
    "        # Denote the state-action value function Q, then the loss of\n",
    "        # prediction error w.r.t. the weights can be separated into two\n",
    "        # parts (the chain rule):\n",
    "        #     dLoss / dweight = (dLoss / dQ) * (dQ / dweight)\n",
    "        # We call the first one loss_grad, and the latter one\n",
    "        # value_grad. We consider the Mean Square Error between the target\n",
    "        # value (G) and the predicted value (Q(s_t, a_t)) to be the loss.\n",
    "        loss_grad = np.zeros((self.act_dim, 1))\n",
    "        # [TODO] fill the propoer value of loss_grad, denoting the gradient\n",
    "        # of the MSE w.r.t. the output of the linear function.\n",
    "        # Hint: only one element of loss_grad is not zero.\n",
    "        loss_grad[actions[tau]] = -(G - self.compute_values(processed_states[tau])[actions[tau]])\n",
    "\n",
    "        # [TODO] compute the value of value_grad, denoting the gradient of\n",
    "        # the output of the linear function w.r.t. the parameters.\n",
    "        value_grad = None\n",
    "        value_grad = np.zeros((self.obs_dim,1))\n",
    "        value_grad[:,0] = processed_states[tau]\n",
    "\n",
    "        assert loss_grad.shape == (self.act_dim, 1)\n",
    "        assert value_grad.shape == (self.obs_dim, 1)\n",
    "\n",
    "        # [TODO] merge two gradients to get the gradient of loss w.r.t. the\n",
    "        # parameters.\n",
    "        gradient = None\n",
    "        gradient = np.dot(value_grad,loss_grad.T)\n",
    "    \n",
    "        return gradient\n",
    "\n",
    "    def apply_gradient(self, gradient):\n",
    "        \"\"\"Apply the gradient to the parameter.\"\"\"\n",
    "        assert gradient.shape == self.parameters.shape, (\n",
    "            gradient.shape, self.parameters.shape)\n",
    "        # [TODO] apply the gradient to self.parameters\n",
    "        self.parameters = self.parameters - self.learning_rate * gradient\n",
    "        pass\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initialize parameters with shape: (4, 2).\n",
      "Now your codes should be bug-free.\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "# Build the test trainer.\n",
    "test_trainer = LinearTrainer(dict(parameter_std=0.0))\n",
    "\n",
    "# Test self.parameters.\n",
    "assert test_trainer.parameters.std() == 0.0, \\\n",
    "    \"Parameters should subjects to a normal distribution with standard \" \\\n",
    "    \"deviation config['parameter_std'], but you have {}.\" \\\n",
    "    \"\".format(test_trainer.parameters.std())\n",
    "assert test_trainer.parameters.mean() == 0, \\\n",
    "    \"Parameters should subjects to a normal distribution with mean 0. \" \\\n",
    "    \"But you have {}.\".format(test_trainer.parameters.mean())\n",
    "\n",
    "# Test compute_values\n",
    "fake_state = test_trainer.env.observation_space.sample()\n",
    "processed_state = test_trainer.process_state(fake_state)\n",
    "assert processed_state.shape == (test_trainer.obs_dim, ), processed_state.shape\n",
    "values = test_trainer.compute_values(fake_state)\n",
    "assert values.shape == (test_trainer.act_dim, ), values.shape\n",
    "\n",
    "# Test compute_gradient\n",
    "tmp_gradient = test_trainer.compute_gradient(\n",
    "    [processed_state]*10, [test_trainer.env.action_space.sample()]*10, [0.0]*10, 2, 5)\n",
    "assert tmp_gradient.shape == test_trainer.parameters.shape\n",
    "\n",
    "test_trainer.train()\n",
    "print(\"Now your codes should be bug-free.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initialize parameters with shape: (4, 2).\n",
      "(0.0s,+0.0s)\tIteration 0, current mean episode reward is 9.28. \n",
      "(0.4s,+0.4s)\tIteration 1000, current mean episode reward is 9.6. \n",
      "(0.7s,+0.3s)\tIteration 2000, current mean episode reward is 9.72. \n",
      "(1.1s,+0.3s)\tIteration 3000, current mean episode reward is 9.78. \n",
      "(1.4s,+0.4s)\tIteration 4000, current mean episode reward is 9.8. \n",
      "(1.8s,+0.4s)\tIteration 5000, current mean episode reward is 9.84. \n",
      "(2.2s,+0.4s)\tIteration 6000, current mean episode reward is 9.84. \n",
      "(2.6s,+0.4s)\tIteration 7000, current mean episode reward is 9.84. \n",
      "(3.0s,+0.4s)\tIteration 8000, current mean episode reward is 9.84. \n",
      "(3.4s,+0.4s)\tIteration 9000, current mean episode reward is 9.82. \n",
      "(3.9s,+0.4s)\tIteration 10000, current mean episode reward is 9.84. \n",
      "(4.2s,+0.4s)\tIteration 11000, current mean episode reward is 9.84. \n",
      "(4.6s,+0.4s)\tIteration 12000, current mean episode reward is 9.84. \n",
      "(5.0s,+0.4s)\tIteration 13000, current mean episode reward is 9.82. \n",
      "(5.4s,+0.4s)\tIteration 14000, current mean episode reward is 9.82. \n",
      "(5.8s,+0.4s)\tIteration 15000, current mean episode reward is 9.82. \n",
      "(7.6s,+1.8s)\tIteration 16000, current mean episode reward is 112.08. \n",
      "(10.4s,+2.7s)\tIteration 17000, current mean episode reward is 110.0. \n",
      "(12.9s,+2.5s)\tIteration 18000, current mean episode reward is 89.46. \n",
      "(15.3s,+2.4s)\tIteration 19000, current mean episode reward is 114.38. \n",
      "(17.7s,+2.5s)\tIteration 20000, current mean episode reward is 95.5. \n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "linear_trainer, _ = run(LinearTrainer, dict(\n",
    "    max_iteration=20000,\n",
    "    evaluate_interval=1000, \n",
    "    parameter_std=0.01,\n",
    "    learning_rate=0.01,\n",
    "    n=3,\n",
    "    env_name=\"CartPole-v0\"\n",
    "))\n",
    "\n",
    "# It's OK to see bad performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average episode reward for your linear agent in CartPole-v0:  100.0\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "# You should see a pop up window which display the movement of the cart and pole.\n",
    "print(\"Average episode reward for your linear agent in CartPole-v0: \",\n",
    "      linear_trainer.evaluate(1, render=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will notice that the linear trainer only has 8 trainable parameters and its performance is quiet bad. In the following section, we will expand the size of parameters and introduce more features as the input to the system so that the system can learn more complex value function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 2.2: Linear Model with Feature Construction\n",
    "\n",
    "(15 / 100 points)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def GenSeq(n,k):\n",
    "    '''\n",
    "    Generate sequence (c1c2..ck) ck\\in{0,1,...,n}\n",
    "    Input:\n",
    "    - n: integer.\n",
    "    - k: integer.\n",
    "    '''\n",
    "    L = []\n",
    "    if k == 0:\n",
    "        return L\n",
    "    if k == 1:\n",
    "        L = ['%s' %i for i in range(n+1)]\n",
    "        return L\n",
    "    LL = GenSeq(n,k-1)\n",
    "    for i in range(n+1):\n",
    "        L = L + ['%s' % i + j for j in LL]\n",
    "    return L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "linear_fc_config = merge_config(dict(\n",
    "    polynomial_order=1,\n",
    "), linear_approximator_config)\n",
    "\n",
    "\n",
    "def polynomial_feature(sequence, order=1):\n",
    "    \"\"\"Construct the order-n polynomial-basis feature of the state.\n",
    "    Refer to Chapter 9.5.1 of the textbook. We expect to get a \n",
    "    vector of length (n+1)^k as the output.\n",
    "    Example:\n",
    "    When the state is [2, 3, 4], the first order polynomial feature\n",
    "    of the state is [\n",
    "        1,\n",
    "        2,\n",
    "        3,\n",
    "        4,\n",
    "        2 * 3 = 6,\n",
    "        2 * 4 = 8,\n",
    "        3 * 4 = 12,\n",
    "        2 * 3 * 4 = 24\n",
    "    ]\n",
    "    It's OK for function polynomial() to return values in different order.\n",
    "    \"\"\"\n",
    "    # [TODO] finish this function.\n",
    "    n = order\n",
    "    k = len(sequence)\n",
    "    features = np.zeros((n+1)**k)\n",
    "    L = GenSeq(n,k)\n",
    "    for i,l in enumerate(L):\n",
    "        features[i] = 1\n",
    "        for j in range(k):\n",
    "            features[i] = features[i] * (sequence[j]**int(l[j]))\n",
    "    return features\n",
    "\n",
    "assert sorted(polynomial_feature([2, 3, 4])) == [1, 2, 3, 4, 6, 8, 12, 24]\n",
    "assert len(polynomial_feature([2, 3, 4], 2)) == 27\n",
    "assert len(polynomial_feature([2, 3, 4], 3)) == 64\n",
    "\n",
    "class LinearTrainerWithFeatureConstruction(LinearTrainer):\n",
    "    \"\"\"In this class, we will expand the dimension of the state.\n",
    "    This procedure is done at self.process_state function.\n",
    "    The modification of self.obs_dim and the shape of parameters\n",
    "    is also needed.\n",
    "    \"\"\"\n",
    "    def __init__(self, config):\n",
    "        config = merge_config(config, linear_fc_config)\n",
    "        # Initialize the abstract class.\n",
    "        super().__init__(config)\n",
    "\n",
    "        self.polynomial_order = self.config[\"polynomial_order\"]\n",
    "\n",
    "        # Expand the size of observation\n",
    "        self.obs_dim = (self.polynomial_order + 1) ** self.obs_dim\n",
    "\n",
    "        # Since we change self.obs_dim, reset the parameters.\n",
    "        self.initialize_parameters()\n",
    "\n",
    "    def process_state(self, state):\n",
    "        \"\"\"Please finish the polynomial function.\"\"\"\n",
    "        processed = polynomial_feature(state, self.polynomial_order)\n",
    "        processed = np.asarray(processed)\n",
    "        assert len(processed) == self.obs_dim, processed.shape\n",
    "        return processed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initialize parameters with shape: (4, 2).\n",
      "Initialize parameters with shape: (16, 2).\n",
      "(0.4s,+0.4s)\tIteration 0, current mean episode reward is 72.78. \n",
      "(4.1s,+3.7s)\tIteration 1000, current mean episode reward is 15.34. \n",
      "(8.4s,+4.3s)\tIteration 2000, current mean episode reward is 53.58. \n",
      "(13.7s,+5.3s)\tIteration 3000, current mean episode reward is 35.98. \n",
      "(18.9s,+5.2s)\tIteration 4000, current mean episode reward is 27.76. \n",
      "(23.8s,+4.9s)\tIteration 5000, current mean episode reward is 79.74. \n",
      "(28.6s,+4.9s)\tIteration 6000, current mean episode reward is 30.3. \n",
      "(33.5s,+4.9s)\tIteration 7000, current mean episode reward is 98.14. \n",
      "(38.7s,+5.2s)\tIteration 8000, current mean episode reward is 106.24. \n",
      "(44.4s,+5.7s)\tIteration 9000, current mean episode reward is 37.0. \n",
      "(55.3s,+10.9s)\tIteration 10000, current mean episode reward is 109.04. \n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "linear_fc_trainer, _ = run(LinearTrainerWithFeatureConstruction, dict(\n",
    "    max_iteration=10000,\n",
    "    evaluate_interval=1000, \n",
    "    parameter_std=0.01,\n",
    "    learning_rate=0.001,\n",
    "    polynomial_order=1,\n",
    "    n=3,\n",
    "    env_name=\"CartPole-v0\"\n",
    "), reward_threshold=195.0)\n",
    "\n",
    "assert linear_fc_trainer.evaluate() > 20.0, \"The best episode reward happening \" \\\n",
    "    \"during training should be greater than the random baseline, that is greather than 20+.\"\n",
    "\n",
    "# This cell should be finished within 10 minitines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average episode reward for your linear agent with feature constructioin in CartPole-v0:  114.0\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "# You should see a pop up window which display the movement of the cart and pole.\n",
    "print(\"Average episode reward for your linear agent with feature constructioin in CartPole-v0: \",\n",
    "      linear_fc_trainer.evaluate(1, render=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3: Multi-layer Perceptron as the approximiator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, you are required to implement a single agent MLP using purely Numpy package. The differences between MLP and linear function are (1). MLP has a hidden layer which increase its representation capacity (2). MLP can leverage activation function after the output of each layer which introduce not linearity.\n",
    "\n",
    "Consider a MLP with one hidden layer containing 100 neurons and activation function `f()`. We call the layer that accepts the state as input and output the activation **hidden layer**, and the layer that accepts the activation as input and produces the values **output layer**. The activation of the hidden layer is:\n",
    "\n",
    "$$a(s_t) = f( W_h^T s_t)$$\n",
    "\n",
    "obvious the activation is a 100-length vector. The output values is:\n",
    "\n",
    "$$Q(s_t) = f(W_o ^ T a(s_t))$$\n",
    "\n",
    "wherein $W_h, W_o$ are the parameters of hidden layer and output layer, respectively. In this section we do not add activation function and hence $f(x) = x$.\n",
    "\n",
    "Moreover, we also introduce the gradient clipping mechanism. In on-policy learning, the norm of gradient is prone to vary drastically, since the output of Q function is unbounded and it can be as large as possible, which leads to exploding gradient issue. Gradient clipping is used to bound the norm of gradient while keeps the direction of gradient vector unchanged. Concretely, the formulation of gradient clipping is:\n",
    "\n",
    "$$g_{clipped} = g_{original} \\cfrac{c}{\\max(c, \\text{norm}(g))}$$\n",
    "\n",
    "wherein $c$ is a hyperparameter which is `config[\"clip_norm\"]` in our implementation. Gradient clipping bounds the gradient norm to $c$ if the norm of original gradient is greater than $c$. You need to implement this mechanism in function `apply_gradient` in the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "# Build the algorithm-specify config.\n",
    "mlp_trainer_config = merge_config(dict(\n",
    "    parameter_std=0.01,\n",
    "    learning_rate=0.01,\n",
    "    hidden_dim=100,\n",
    "    n=3,\n",
    "    clip_norm=1.0,\n",
    "    clip_gradient=True\n",
    "), default_config)\n",
    "\n",
    "\n",
    "class MLPTrainer(LinearTrainer):\n",
    "    def __init__(self, config):\n",
    "        config = merge_config(config, mlp_trainer_config)\n",
    "        self.hidden_dim = config[\"hidden_dim\"]\n",
    "        super().__init__(config)\n",
    "\n",
    "    def initialize_parameters(self):\n",
    "        # [TODO] Initialize self.hidden_parameters and self.output_parameters,\n",
    "        #  which are two dimensional matrices, and subject to normal\n",
    "        #  distributions with scale config[\"parameter_std\"]\n",
    "        std = self.config[\"parameter_std\"]\n",
    "        self.hidden_parameters = std * np.random.randn(self.obs_dim,self.hidden_dim)\n",
    "        self.output_parameters = std * np.random.randn(self.hidden_dim,self.act_dim)\n",
    "\n",
    "    def compute_values(self, processed_state):\n",
    "        \"\"\"Compute the value for each potential action. Note that you\n",
    "        should NOT preprocess the state here.\"\"\"\n",
    "        assert processed_state.ndim == 1, processed_state.shape\n",
    "        activation = self.compute_activation(processed_state)\n",
    "        \n",
    "        activation = self.compute_activation(np.dot(self.hidden_parameters.T,processed_state))\n",
    "        values = self.compute_activation(np.dot(self.output_parameters.T,activation))\n",
    "        return values\n",
    "\n",
    "    def compute_activation(self, processed_state):\n",
    "        \"\"\"Given a processed state, first we need to compute the activtaion\n",
    "        (the output of hidden layer). Then we compute the values (the output of\n",
    "        the output layer).\n",
    "        \"\"\"\n",
    "        activation = processed_state\n",
    "        return activation\n",
    "\n",
    "    def compute_gradient(self, processed_states, actions, rewards, tau, T):\n",
    "        n = self.n\n",
    "        \n",
    "        # [TODO] compute the target value.\n",
    "        # Hint: copy your codes in LinearTrainer.\n",
    "        G = None\n",
    "        G = 0\n",
    "        for i in range(tau+1,min(tau+n+1,T+1)):\n",
    "            G = G + (self.gamma**(i-tau-1)) * rewards[i]\n",
    "\n",
    "        if tau + n < T:\n",
    "            # Hint: Since we use Sarsa algorithm here,\n",
    "            #  the Q value of time tau+n is the Q value of action\n",
    "            #  in time tau+n. So you should take the tau+n element of\n",
    "            #  processed_states as input to compute the Q values\n",
    "            #  and then take the \"actions[tau+n]\" as the index to get\n",
    "            #  the Q value in tau+n.\n",
    "            Q_tau_plus_n = self.compute_values(processed_states[tau+n])[actions[tau+n]]\n",
    "            \n",
    "            G = G + (self.gamma ** n) * Q_tau_plus_n\n",
    "\n",
    "        # Denote the state-action value function Q, then the loss of\n",
    "        # prediction error w.r.t. the output layer weights can be \n",
    "        # separated into two parts (the chain rule):\n",
    "        #     dError / dweight = (dError / dQ) * (dQ / dweight)\n",
    "        # We call the first one loss_grad, and the latter one\n",
    "        # value_grad. We consider the Mean Square Error between the target\n",
    "        # value (G) and the predict value (Q(s_t, a_t)) to be the loss.\n",
    "        cur_state = processed_states[tau]\n",
    "\n",
    "        loss_grad = np.zeros((self.act_dim, 1))  # [act_dim, 1]\n",
    "        # [TODO] compute loss_grad\n",
    "        loss_grad[actions[tau]] = -(G - self.compute_values(cur_state)[actions[tau]])\n",
    "        \n",
    "        # [TODO] compute the gradient of output layer parameters\n",
    "        # bugs here\n",
    "        output_gradient = None\n",
    "        output_gradient = self.compute_activation(np.dot(self.hidden_parameters.T,cur_state)).reshape(-1,1).dot(loss_grad.T)\n",
    "        \n",
    "        # [TODO] compute the gradient of hidden layer parameters\n",
    "        # Hint: using chain rule and derive the formulation\n",
    "        hidden_gradient = None\n",
    "        hidden_gradient = np.dot(cur_state.reshape(-1,1),np.dot(self.output_parameters,loss_grad).T)\n",
    "    \n",
    "        assert np.all(np.isfinite(output_gradient)), \\\n",
    "            \"Invalid value occurs in output_gradient! {}\".format(\n",
    "                output_gradient)\n",
    "        assert np.all(np.isfinite(hidden_gradient)), \\\n",
    "            \"Invalid value occurs in hidden_gradient! {}\".format(\n",
    "                hidden_gradient)\n",
    "        return [hidden_gradient, output_gradient]\n",
    "\n",
    "    def apply_gradient(self, gradients):\n",
    "        \"\"\"Apply the gradientss to the two layers' parameters.\"\"\"\n",
    "        assert len(gradients) == 2\n",
    "        hidden_gradient, output_gradient = gradients\n",
    "\n",
    "        assert output_gradient.shape == (self.hidden_dim, self.act_dim)\n",
    "        assert hidden_gradient.shape == (self.obs_dim, self.hidden_dim)\n",
    "        \n",
    "        # [TODO] Implement the clip gradient mechansim\n",
    "        # Hint: when the old gradient has norm less that clip_norm,\n",
    "        #  then nothing happens. Otherwise shrink the gradient to\n",
    "        #  make its norm equal to clip_norm.\n",
    "        if self.config[\"clip_gradient\"]:\n",
    "            clip_norm = self.config[\"clip_norm\"]\n",
    "            hidden_gradient = hidden_gradient * clip_norm / max(clip_norm, np.linalg.norm(hidden_gradient))\n",
    "            output_gradient = output_gradient * clip_norm / max(clip_norm, np.linalg.norm(output_gradient))\n",
    "\n",
    "        # [TODO] update the parameters\n",
    "        # Hint: Remember to check the sign when applying the gradient\n",
    "        #  into the parameters. Should you add or minus the gradients?\n",
    "        self.hidden_parameters = self.hidden_parameters - self.learning_rate * hidden_gradient\n",
    "        self.output_parameters = self.output_parameters - self.learning_rate * output_gradient\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Now let's see what happen if clip gradient is not enable!\n",
      "(0.4s,+0.4s)\tIteration 0, current mean episode reward is 33.52. \n",
      "(0.9s,+0.4s)\tIteration 100, current mean episode reward is 92.24. \n",
      "(1.3s,+0.5s)\tIteration 200, current mean episode reward is 87.3. \n",
      "(1.7s,+0.4s)\tIteration 300, current mean episode reward is 79.46. \n",
      "(2.2s,+0.5s)\tIteration 400, current mean episode reward is 79.7. \n",
      "(2.6s,+0.4s)\tIteration 500, current mean episode reward is 73.84. \n",
      "(3.1s,+0.4s)\tIteration 600, current mean episode reward is 76.54. \n",
      "(3.4s,+0.4s)\tIteration 700, current mean episode reward is 63.74. \n",
      "(3.8s,+0.4s)\tIteration 800, current mean episode reward is 60.8. \n",
      "(4.2s,+0.4s)\tIteration 900, current mean episode reward is 59.18. \n",
      "(4.5s,+0.3s)\tIteration 1000, current mean episode reward is 56.48. \n",
      "(4.8s,+0.3s)\tIteration 1100, current mean episode reward is 52.9. \n",
      "(5.1s,+0.3s)\tIteration 1200, current mean episode reward is 48.04. \n",
      "(5.4s,+0.3s)\tIteration 1300, current mean episode reward is 46.54. \n",
      "(5.7s,+0.3s)\tIteration 1400, current mean episode reward is 42.66. \n",
      "(5.9s,+0.3s)\tIteration 1500, current mean episode reward is 45.02. \n",
      "(6.2s,+0.3s)\tIteration 1600, current mean episode reward is 40.32. \n",
      "(6.4s,+0.3s)\tIteration 1700, current mean episode reward is 44.48. \n",
      "(6.7s,+0.3s)\tIteration 1800, current mean episode reward is 42.72. \n",
      "(7.0s,+0.3s)\tIteration 1900, current mean episode reward is 40.04. \n",
      "(7.2s,+0.2s)\tIteration 2000, current mean episode reward is 33.58. \n",
      "(7.4s,+0.2s)\tIteration 2100, current mean episode reward is 33.58. \n",
      "(7.6s,+0.2s)\tIteration 2200, current mean episode reward is 32.98. \n",
      "(7.8s,+0.2s)\tIteration 2300, current mean episode reward is 33.58. \n",
      "(8.0s,+0.2s)\tIteration 2400, current mean episode reward is 32.94. \n",
      "(8.2s,+0.2s)\tIteration 2500, current mean episode reward is 32.66. \n",
      "(8.4s,+0.2s)\tIteration 2600, current mean episode reward is 33.08. \n",
      "(8.6s,+0.2s)\tIteration 2700, current mean episode reward is 33.22. \n",
      "(8.9s,+0.2s)\tIteration 2800, current mean episode reward is 33.28. \n",
      "(9.1s,+0.2s)\tIteration 2900, current mean episode reward is 33.1. \n",
      "(9.3s,+0.2s)\tIteration 3000, current mean episode reward is 33.24. \n",
      "We expect to see bad performance (<195). The performance without gradient clipping: 33.24.\n",
      "Try next cell to see the impact of gradient clipping.\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "print(\"Now let's see what happen if clip gradient is not enable!\")\n",
    "try:\n",
    "    failed_mlp_trainer, _ = run(MLPTrainer, dict(\n",
    "        max_iteration=3000,\n",
    "        evaluate_interval=100, \n",
    "        parameter_std=0.01,\n",
    "        learning_rate=0.001,\n",
    "        hidden_dim=100,\n",
    "        clip_gradient=False,  # <<< Gradient clipping is OFF!\n",
    "        env_name=\"CartPole-v0\"\n",
    "    ), reward_threshold=195.0)\n",
    "    print(\"We expect to see bad performance (<195). \"\n",
    "          \"The performance without gradient clipping: {}.\"\n",
    "          \"\".format(failed_mlp_trainer.evaluate()))\n",
    "except AssertionError as e:\n",
    "    print(traceback.format_exc())\n",
    "    print(\"Infinity happen during training. It's OK since the gradient is not bounded.\")\n",
    "finally:\n",
    "    print(\"Try next cell to see the impact of gradient clipping.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Now let's see what happen if clip gradient is not enable!\n",
      "(0.5s,+0.5s)\tIteration 0, current mean episode reward is 33.52. \n",
      "(1.0s,+0.6s)\tIteration 100, current mean episode reward is 92.44. \n",
      "(1.6s,+0.6s)\tIteration 200, current mean episode reward is 93.76. \n",
      "(2.3s,+0.6s)\tIteration 300, current mean episode reward is 98.18. \n",
      "(2.9s,+0.7s)\tIteration 400, current mean episode reward is 103.18. \n",
      "(3.6s,+0.7s)\tIteration 500, current mean episode reward is 108.46. \n",
      "(4.3s,+0.7s)\tIteration 600, current mean episode reward is 111.26. \n",
      "(5.0s,+0.7s)\tIteration 700, current mean episode reward is 121.48. \n",
      "(5.7s,+0.7s)\tIteration 800, current mean episode reward is 135.22. \n",
      "(6.6s,+0.9s)\tIteration 900, current mean episode reward is 178.46. \n",
      "(7.8s,+1.1s)\tIteration 1000, current mean episode reward is 198.96. \n",
      "In 1000 iteration, current mean episode reward 198.960 is greater than reward threshold 195.0. Congratulation! Now we exit the training process.\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "print(\"Now let's see what happen if clip gradient is not enable!\")\n",
    "mlp_trainer, _ = run(MLPTrainer, dict(\n",
    "    max_iteration=3000,\n",
    "    evaluate_interval=100, \n",
    "    parameter_std=0.01,\n",
    "    learning_rate=0.001,\n",
    "    hidden_dim=100,\n",
    "    clip_gradient=True,  # <<< Gradient clipping is ON!\n",
    "    env_name=\"CartPole-v0\"\n",
    "), reward_threshold=195.0)\n",
    "\n",
    "assert mlp_trainer.evaluate() > 195.0, \"Check your codes. \" \\\n",
    "    \"Your agent should achieve {} reward in 200 iterations.\" \\\n",
    "    \"But it achieve {} reward in evaluation.\"\n",
    "\n",
    "# In our implementation, the task is solved in 200 iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average episode reward for your MLP agent with gradient clipping in CartPole-v0:  200.0\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "# You should see a pop up window which display the movement of the cart and pole.\n",
    "print(\"Average episode reward for your MLP agent with gradient clipping in CartPole-v0: \",\n",
    "      mlp_trainer.evaluate(1, render=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Interesting right? The gradient clipping technique makes the training converge much faster!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 4: Implement Deep Q Learning in Pytorch\n",
    "\n",
    "(50 / 100 points)\n",
    "\n",
    "In this section, you will get familiar with the basic logic of pytorch, which lay the ground for further learning. We will implement a MLP similar to the one in Section 3 using Pytorch, a powerful Deep Learning framework. Before start, you need to make sure using `pip install torch` to install it.\n",
    "\n",
    "If you are not familiar with Pytorch, we suggest you to go through pytorch official quickstart tutorials:\n",
    "1. [quickstart](https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html)\n",
    "2. [tutorial on RL](https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html)\n",
    "\n",
    "Different from the algorithm in Section 3, we will implement Deep Q Network (DQN) in this section. The main differences are concluded as following:\n",
    "\n",
    "**DQN requires an experience replay memory to store the transitions.** A replay memory is implemented in the following `ExperienceReplayMemory` class. It can contain a certain amount of transitions: `(s_t, a_t, r_t, s_t+1, done_t)`. When the memory is full, the earliest transition is discarded to store the latest one.\n",
    "\n",
    "The introduction of replay memory increase the sample efficiency (since each transition might be used multiple times) when solving complex task, though you may find it learn slowly in this assignment since the CartPole-v0 is a relatively easy environment.\n",
    "\n",
    "\n",
    "**DQN is an off-policy algorithm and has difference when computing TD error, compared to Sarsa.** In Sarsa, the TD error is computed as: \n",
    "\n",
    "$$(r_t + \\gamma Q(s_{t+1}, a_{t+1}) - Q(s_t, a_t))$$ \n",
    "\n",
    "wherein the next action $a_{t+1}$ is the one the policy selects. However, in traditional Q learning, it assume the next action is the one that maximizes the action values and use this assumption to compute the TD: \n",
    "\n",
    "$$(r_t + \\gamma \\max_{a_{t+1}} Q(s_{t+1}, a_{t+1}) - Q(s_t, a_t))$$\n",
    "\n",
    "**DQN has make delayed update target network, which is another difference even compared to the traditional Q learning.** DQN maintains another neural network called target network that has identical structure of the Q network. After a certain amount of steps has been taken, the target network copies the parameters of the Q network to itself. Normally, the update of target network is much less frequent than the update of the Q network. The Q network is updated in each step.\n",
    "\n",
    "The reason to leverage the target network is to stabilize the estimation of TD error. In DQN, the TD error is evaluated as:\n",
    "\n",
    "$$(r_t + \\gamma \\max_{a_{t+1}} Q^{target}(s_{t+1}, a_{t+1}) - Q(s_t, a_t))$$\n",
    "\n",
    "The Q values of next state is estimated by the target network, not the Q network that is updating. This mechanism can reduce the variance of gradient because the estimation of Q values of next states is not influenced by the update of the Q network.\n",
    "\n",
    "In the engineering aspect, the differences between `DQNTrainer` and the previous `MLPTrainer` are:\n",
    "\n",
    "1. DQN uses pytorch model to serve as the approximator. So we need to rewrite the `initialize_parameter` function to build the pytorch model. Also the `train` function is changed since the gradient optimization is conducted by pytorch, therefore we need to write the pytorch pipeline in `train`.\n",
    "2. DQN has replay memory. So we need to initialize it, feed data into it and take the transitions out.\n",
    "3. Thank to the replay memory and pytorch, DQN can be updated in a batch. So you need to carefully compute the Q target via matrix computation.\n",
    "4. We use Adam optimizer to conduct the gradient optimization. You need to get familiar with how to compute the loss and conduct backward propagation.\n",
    "\n",
    "\\* Note that in this assignment, we use purely CPU to train the agent. You do not need to right your pytorch code that compatible with GPU and also do not need to run your code in a GPU machine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "from collections import deque\n",
    "import random\n",
    "\n",
    "class ExperienceReplayMemory:\n",
    "    \"\"\"Store and sample the transitions\"\"\"\n",
    "    def __init__(self, capacity):\n",
    "        # deque is a useful class which acts like a list but only contain\n",
    "        # finite elements.When appending new element make deque exceeds the \n",
    "        # `maxlen`, the oldest element (the index 0 element) will be removed.\n",
    "        \n",
    "        # [TODO] uncomment next line. \n",
    "        self.memory = deque(maxlen=capacity)\n",
    "\n",
    "    def push(self, transition):\n",
    "        self.memory.append(transition)\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        return random.sample(self.memory, batch_size)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "class PytorchModel(nn.Module):\n",
    "    def __init__(self, obs_dim, act_dim):\n",
    "        super(PytorchModel, self).__init__()\n",
    "        \n",
    "        # [TODO] Build a sequential model with two hidden layers.\n",
    "        # The first hidden layer has 100 hidden units, followed by\n",
    "        # a ReLU activation function.\n",
    "        # The second output layer take the activation vector as input\n",
    "        # and pass through another 100 hidden units, also followed\n",
    "        # by a ReLU activation function.\n",
    "        # Then the output 100-length vector passes the output layer\n",
    "        # and output the values of all actions.\n",
    "        # \n",
    "        #     input (length: obs_dim) \n",
    "        # ->  100 units hidden layer\n",
    "        # ->  ReLU \n",
    "        # ->  100 units hidden layer\n",
    "        # ->  ReLU\n",
    "        # ->  act_dim units hidden layer\n",
    "        # -> output (length: act_dim)\n",
    "        self.action_value = nn.Sequential(\n",
    "            nn.Linear(obs_dim,100),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(100,100),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(100,act_dim)\n",
    "        )\n",
    "\n",
    "    def forward(self, obs):\n",
    "        return self.action_value(obs)\n",
    "    \n",
    "# Test\n",
    "test_model = PytorchModel(3, 7)\n",
    "assert isinstance(test_model.action_value, nn.Module)\n",
    "assert test_model(torch.from_numpy(np.ones([3], dtype=np.float32))).shape \\\n",
    "    == torch.Size([7])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "pytorch_config = merge_config(dict(\n",
    "    memory_size=50000,\n",
    "    learn_start=5000,\n",
    "    batch_size=32,\n",
    "    target_update_freq=500,  # in steps\n",
    "    learn_freq=1,  # in steps\n",
    "    n=1\n",
    "), mlp_trainer_config)\n",
    "\n",
    "\n",
    "def to_tensor(x):\n",
    "    \"\"\"A helper function to transform a numpy array to a Pytorch Tensor\"\"\"\n",
    "    if isinstance(x, np.ndarray):\n",
    "        x = torch.from_numpy(x).type(torch.float32)\n",
    "    assert isinstance(x, torch.Tensor)\n",
    "    if x.dim() == 3 or x.dim() == 1:\n",
    "        x = x.unsqueeze(0)\n",
    "    assert x.dim() == 2 or x.dim() == 4, x.shape\n",
    "    return x\n",
    "\n",
    "\n",
    "class DQNTrainer(MLPTrainer):\n",
    "    def __init__(self, config):\n",
    "        config = merge_config(config, pytorch_config)\n",
    "        self.learning_rate = config[\"learning_rate\"]\n",
    "        super().__init__(config)\n",
    "\n",
    "        self.memory = ExperienceReplayMemory(config[\"memory_size\"])\n",
    "        self.learn_start = config[\"learn_start\"]\n",
    "        self.batch_size = config[\"batch_size\"]\n",
    "        self.target_update_freq = config[\"target_update_freq\"]\n",
    "        self.clip_norm = config[\"clip_norm\"]\n",
    "        self.step_since_update = 0\n",
    "        self.total_step = 0\n",
    "\n",
    "    def initialize_parameters(self):\n",
    "        \"\"\"Initialize the pytorch model as the Q network and the target network\"\"\"\n",
    "        # [TODO] Initialize two network using PytorchModel class\n",
    "        self.network = PytorchModel(self.obs_dim,self.act_dim)\n",
    "\n",
    "        self.network.eval()\n",
    "        self.network.share_memory()\n",
    "\n",
    "        # [TODO] Initialize target network, which is identical to self.network,\n",
    "        # and should have the same weights with self.network. So you should\n",
    "        # put the weights of self.network into self.target_network.\n",
    "        self.target_network = PytorchModel(self.obs_dim,self.act_dim)\n",
    "        self.target_network.load_state_dict(self.network.state_dict())\n",
    "\n",
    "        self.target_network.eval()\n",
    "\n",
    "        # Build Adam optimizer and MSE Loss.\n",
    "        # [TODO] Uncomment next few lines\n",
    "        self.optimizer = torch.optim.Adam(\n",
    "            self.network.parameters(), lr=self.learning_rate\n",
    "        )\n",
    "        self.loss = nn.MSELoss()\n",
    "\n",
    "    def compute_values(self, processed_state):\n",
    "        \"\"\"Compute the value for each potential action. Note that you\n",
    "        should NOT preprocess the state here.\"\"\"\n",
    "        # [TODO] Convert the output of neural network to numpy array\n",
    "        values = self.network.forward(processed_state).detach().numpy()\n",
    "        return values\n",
    "\n",
    "    def train(self):\n",
    "        s = self.env.reset()\n",
    "        processed_s = self.process_state(s)\n",
    "        act = self.compute_action(processed_s)\n",
    "        stat = {\"loss\": []}\n",
    "\n",
    "        for t in range(self.max_episode_length):\n",
    "            next_state, reward, done, _ = self.env.step(act)\n",
    "            next_processed_s = self.process_state(next_state)\n",
    "\n",
    "            # Push the transition into memory.\n",
    "            self.memory.push(\n",
    "                (processed_s, act, reward, next_processed_s, done)\n",
    "            )\n",
    "\n",
    "            processed_s = next_processed_s\n",
    "            act = self.compute_action(next_processed_s)\n",
    "            self.step_since_update += 1\n",
    "            self.total_step += 1\n",
    "\n",
    "            if done:\n",
    "                break\n",
    "                \n",
    "            if t % self.config[\"learn_freq\"] != 0:\n",
    "                # It's not necessary to update in each step.\n",
    "                continue\n",
    "\n",
    "            if len(self.memory) < self.learn_start:\n",
    "                continue\n",
    "            elif len(self.memory) == self.learn_start:\n",
    "                print(\"Current memory contains {} transitions, \"\n",
    "                      \"start learning!\".format(self.learn_start))\n",
    "\n",
    "            batch = self.memory.sample(self.batch_size)\n",
    "\n",
    "            # Transform a batch of state / action / .. into a tensor.\n",
    "            state_batch = to_tensor(\n",
    "                np.stack([transition[0] for transition in batch])\n",
    "            )\n",
    "            action_batch = to_tensor(\n",
    "                np.stack([transition[1] for transition in batch])\n",
    "            )\n",
    "            reward_batch = to_tensor(\n",
    "                np.stack([transition[2] for transition in batch])\n",
    "            )\n",
    "            next_state_batch = torch.stack(\n",
    "                [transition[3] for transition in batch]\n",
    "            )\n",
    "            done_batch = to_tensor(\n",
    "                np.stack([transition[4] for transition in batch])\n",
    "            )\n",
    "\n",
    "            with torch.no_grad():\n",
    "                # [TODO] Compute the values of Q in next state in batch.\n",
    "                # Hint: \n",
    "                #  1. Q_t_plus_one is the maximum value of Q values of possible\n",
    "                #     actions in next state. So the input to the network is \n",
    "                #     next_state_batch.\n",
    "                #  2. Q_t_plus_one is computed using the target network.\n",
    "                values = self.target_network.forward(next_state_batch).detach().numpy() \n",
    "                Q_t_plus_one = torch.from_numpy(np.max(values,axis=1))\n",
    "                \n",
    "                assert isinstance(Q_t_plus_one, torch.Tensor)\n",
    "                assert Q_t_plus_one.dim() == 1\n",
    "                \n",
    "                # [TODO] Compute the target value of Q in batch.\n",
    "                # Hint: The Q target is simply r_t + gamma * Q_t+1 \n",
    "                #  IF the episode is not done at time t.\n",
    "                #  That is, the (gamma*Q_t+1) term should be masked out\n",
    "                #  if done_batch[t] is True.\n",
    "                #  A smart way to do so is: using (1-done_batch) as multiplier\n",
    "                Q_target = torch.zeros((self.batch_size,))\n",
    "                \n",
    "                for i in range(self.batch_size):\n",
    "                    Q_target[i] = reward_batch[0,i] + self.gamma * Q_t_plus_one[i] * (1 - done_batch[0,i])\n",
    "                \n",
    "                assert Q_target.shape == (self.batch_size,)\n",
    "            \n",
    "            # [TODO] Collect the Q values in batch.\n",
    "            # Hint: Remember to call self.network.train()\n",
    "            #  before you get the Q value from self.network(state_batch),\n",
    "            #  otherwise the graident will not be recorded by pytorch.\n",
    "            self.network.train()\n",
    "            index = action_batch.long().reshape((self.batch_size,1))\n",
    "            v = self.network(state_batch) \n",
    "            Q_t = torch.gather(v,1,index).squeeze()\n",
    "            assert Q_t.shape == Q_target.shape\n",
    "\n",
    "            # Update the network\n",
    "            self.optimizer.zero_grad()\n",
    "            loss = self.loss(input=Q_t, target=Q_target)\n",
    "            loss_value = loss.item()\n",
    "            stat['loss'].append(loss_value)\n",
    "            loss.backward()\n",
    "            \n",
    "            # [TODO] Gradient clipping. Uncomment next line\n",
    "            nn.utils.clip_grad_norm_(self.network.parameters(), self.clip_norm)\n",
    "            \n",
    "            self.optimizer.step()\n",
    "            self.network.eval()\n",
    "\n",
    "        if len(self.memory) >= self.learn_start and \\\n",
    "                self.step_since_update > self.target_update_freq:\n",
    "            print(\"{} steps has passed since last update. Now update the\"\n",
    "                  \" parameter of the behavior policy. Current step: {}\".format(\n",
    "                self.step_since_update, self.total_step\n",
    "            ))\n",
    "            self.step_since_update = 0\n",
    "            # [TODO] Copy the weights of self.network to self.target_network.\n",
    "            self.target_network.load_state_dict(self.network.state_dict())\n",
    "            \n",
    "            self.target_network.eval()\n",
    "            \n",
    "        return {\"loss\": np.mean(stat[\"loss\"]), \"episode_len\": t}\n",
    "\n",
    "    def process_state(self, state):\n",
    "        return torch.from_numpy(state).type(torch.float32)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Now your codes should be bug-free.\n",
      "(0.1s,+0.1s)\tIteration 0, current mean episode reward is 9.24. {'loss': nan, 'episode_len': 9.0}\n",
      "Current memory contains 100 transitions, start learning!\n",
      "(0.3s,+0.2s)\tIteration 10, current mean episode reward is 9.24. {'loss': 0.0303, 'episode_len': 10.0}\n",
      "(1.0s,+0.7s)\tIteration 20, current mean episode reward is 9.9. {'loss': 0.0011, 'episode_len': 7.0}\n",
      "Test passed!\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "# Build the test trainer.\n",
    "test_trainer = DQNTrainer({})\n",
    "\n",
    "# Test compute_values\n",
    "fake_state = test_trainer.env.observation_space.sample()\n",
    "processed_state = test_trainer.process_state(fake_state)\n",
    "assert processed_state.shape == (test_trainer.obs_dim, ), processed_state.shape\n",
    "values = test_trainer.compute_values(processed_state)\n",
    "assert values.shape == (test_trainer.act_dim, ), values.shape\n",
    "\n",
    "test_trainer.train()\n",
    "print(\"Now your codes should be bug-free.\")\n",
    "\n",
    "_ = run(DQNTrainer, dict(\n",
    "    max_iteration=20,\n",
    "    evaluate_interval=10, \n",
    "    learn_start=100,\n",
    "    env_name=\"CartPole-v0\",\n",
    "))\n",
    "\n",
    "print(\"Test passed!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0.1s,+0.1s)\tIteration 0, current mean episode reward is 9.48. {'loss': nan, 'episode_len': 8.0}\n",
      "(0.2s,+0.1s)\tIteration 10, current mean episode reward is 9.48. {'loss': nan, 'episode_len': 10.0}\n",
      "(0.3s,+0.1s)\tIteration 20, current mean episode reward is 9.48. {'loss': nan, 'episode_len': 10.0}\n",
      "(0.5s,+0.1s)\tIteration 30, current mean episode reward is 9.48. {'loss': nan, 'episode_len': 8.0}\n",
      "(0.6s,+0.1s)\tIteration 40, current mean episode reward is 9.48. {'loss': nan, 'episode_len': 10.0}\n",
      "(0.7s,+0.1s)\tIteration 50, current mean episode reward is 9.48. {'loss': nan, 'episode_len': 8.0}\n",
      "(0.8s,+0.1s)\tIteration 60, current mean episode reward is 9.48. {'loss': nan, 'episode_len': 9.0}\n",
      "(1.0s,+0.1s)\tIteration 70, current mean episode reward is 9.48. {'loss': nan, 'episode_len': 9.0}\n",
      "(1.1s,+0.1s)\tIteration 80, current mean episode reward is 9.48. {'loss': nan, 'episode_len': 10.0}\n",
      "(1.2s,+0.1s)\tIteration 90, current mean episode reward is 9.48. {'loss': nan, 'episode_len': 9.0}\n",
      "(1.3s,+0.1s)\tIteration 100, current mean episode reward is 9.48. {'loss': nan, 'episode_len': 9.0}\n",
      "Current memory contains 1000 transitions, start learning!\n",
      "1004 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 1004\n",
      "(2.0s,+0.7s)\tIteration 110, current mean episode reward is 9.24. {'loss': 0.214, 'episode_len': 38.0}\n",
      "(2.6s,+0.6s)\tIteration 120, current mean episode reward is 9.48. {'loss': 0.1787, 'episode_len': 7.0}\n",
      "(3.1s,+0.6s)\tIteration 130, current mean episode reward is 9.24. {'loss': 0.18, 'episode_len': 17.0}\n",
      "(3.6s,+0.5s)\tIteration 140, current mean episode reward is 9.24. {'loss': 0.0997, 'episode_len': 9.0}\n",
      "(4.2s,+0.5s)\tIteration 150, current mean episode reward is 9.28. {'loss': 0.1373, 'episode_len': 9.0}\n",
      "(4.7s,+0.5s)\tIteration 160, current mean episode reward is 9.48. {'loss': 0.1692, 'episode_len': 8.0}\n",
      "(5.2s,+0.5s)\tIteration 170, current mean episode reward is 9.24. {'loss': 0.1908, 'episode_len': 8.0}\n",
      "(5.7s,+0.6s)\tIteration 180, current mean episode reward is 9.24. {'loss': 0.1405, 'episode_len': 13.0}\n",
      "(6.3s,+0.6s)\tIteration 190, current mean episode reward is 9.24. {'loss': 0.0897, 'episode_len': 11.0}\n",
      "1003 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 2007\n",
      "(6.9s,+0.5s)\tIteration 200, current mean episode reward is 9.42. {'loss': 0.0364, 'episode_len': 9.0}\n",
      "(7.4s,+0.6s)\tIteration 210, current mean episode reward is 9.76. {'loss': 0.0212, 'episode_len': 9.0}\n",
      "(8.0s,+0.6s)\tIteration 220, current mean episode reward is 9.6. {'loss': 0.0297, 'episode_len': 9.0}\n",
      "(8.6s,+0.5s)\tIteration 230, current mean episode reward is 9.7. {'loss': 0.0376, 'episode_len': 12.0}\n",
      "(9.1s,+0.6s)\tIteration 240, current mean episode reward is 10.06. {'loss': 0.0214, 'episode_len': 9.0}\n",
      "(9.7s,+0.6s)\tIteration 250, current mean episode reward is 9.88. {'loss': 0.0309, 'episode_len': 11.0}\n",
      "(10.3s,+0.6s)\tIteration 260, current mean episode reward is 9.92. {'loss': 0.0349, 'episode_len': 10.0}\n",
      "(10.9s,+0.6s)\tIteration 270, current mean episode reward is 9.98. {'loss': 0.0239, 'episode_len': 8.0}\n",
      "(11.4s,+0.5s)\tIteration 280, current mean episode reward is 9.9. {'loss': 0.0282, 'episode_len': 9.0}\n",
      "(12.0s,+0.6s)\tIteration 290, current mean episode reward is 9.68. {'loss': 0.0441, 'episode_len': 10.0}\n",
      "1004 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 3011\n",
      "(12.6s,+0.6s)\tIteration 300, current mean episode reward is 10.44. {'loss': 0.0312, 'episode_len': 13.0}\n",
      "(13.3s,+0.6s)\tIteration 310, current mean episode reward is 10.42. {'loss': 0.0194, 'episode_len': 10.0}\n",
      "(13.9s,+0.6s)\tIteration 320, current mean episode reward is 11.62. {'loss': 0.0323, 'episode_len': 8.0}\n",
      "(14.5s,+0.6s)\tIteration 330, current mean episode reward is 12.56. {'loss': 0.0227, 'episode_len': 9.0}\n",
      "(15.1s,+0.6s)\tIteration 340, current mean episode reward is 13.32. {'loss': 0.0323, 'episode_len': 9.0}\n",
      "(15.8s,+0.7s)\tIteration 350, current mean episode reward is 10.4. {'loss': 0.0358, 'episode_len': 11.0}\n",
      "(16.6s,+0.8s)\tIteration 360, current mean episode reward is 10.54. {'loss': 0.0535, 'episode_len': 14.0}\n",
      "(17.3s,+0.7s)\tIteration 370, current mean episode reward is 11.26. {'loss': 0.0248, 'episode_len': 9.0}\n",
      "1004 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 4015\n",
      "(19.1s,+1.8s)\tIteration 380, current mean episode reward is 43.98. {'loss': 0.0467, 'episode_len': 11.0}\n",
      "(21.0s,+1.9s)\tIteration 390, current mean episode reward is 51.72. {'loss': 0.0302, 'episode_len': 13.0}\n",
      "(23.3s,+2.3s)\tIteration 400, current mean episode reward is 59.26. {'loss': 0.0812, 'episode_len': 10.0}\n",
      "1006 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 5021\n",
      "(25.2s,+2.0s)\tIteration 410, current mean episode reward is 77.74. {'loss': 0.0975, 'episode_len': 65.0}\n",
      "(28.1s,+2.9s)\tIteration 420, current mean episode reward is 36.2. {'loss': 0.1309, 'episode_len': 19.0}\n",
      "1028 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 6049\n",
      "(32.8s,+4.6s)\tIteration 430, current mean episode reward is 191.08. {'loss': 0.1676, 'episode_len': 199.0}\n",
      "(38.8s,+6.0s)\tIteration 440, current mean episode reward is 200.0. {'loss': 0.0698, 'episode_len': 26.0}\n",
      "In 440 iteration, current mean episode reward 200.000 is greater than reward threshold 195.0. Congratulation! Now we exit the training process.\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "pytorch_trainer, pytorch_stat = run(DQNTrainer, dict(\n",
    "    max_iteration=2000,\n",
    "    evaluate_interval=10, \n",
    "    learning_rate=0.01,\n",
    "    clip_norm=10.0,\n",
    "    memory_size=50000,\n",
    "    learn_start=1000,\n",
    "    eps=0.1,\n",
    "    target_update_freq=1000,\n",
    "    batch_size=32,\n",
    "    env_name=\"CartPole-v0\",\n",
    "), reward_threshold=195.0)\n",
    "\n",
    "reward = pytorch_trainer.evaluate()\n",
    "assert reward > 195.0, \"Check your codes. \" \\\n",
    "    \"Your agent should achieve {} reward in 1000 iterations.\" \\\n",
    "    \"But it achieve {} reward in evaluation.\".format(195.0, reward)\n",
    "\n",
    "# Should solve the task in 10 minutes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average episode reward for your Pytorch agent in CartPole-v0:  200.0\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "# You should see a pop up window which display the movement of the cart and pole.\n",
    "print(\"Average episode reward for your Pytorch agent in CartPole-v0: \",\n",
    "      pytorch_trainer.evaluate(1, render=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(7.4s,+7.4s)\tIteration 0, current mean episode reward is -21.0. {'loss': nan, 'episode_len': 1022.0}\n",
      "10194 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 10194\n",
      "1019 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 11213\n",
      "(22.5s,+15.1s)\tIteration 10, current mean episode reward is -21.0. {'loss': 7.4894, 'episode_len': 1018.0}\n",
      "1043 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 12256\n",
      "1029 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 13285\n",
      "1009 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 14294\n",
      "1226 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 15520\n",
      "1106 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 16626\n",
      "1033 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 17659\n",
      "1047 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 18706\n",
      "1101 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 19807\n",
      "1253 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 21060\n",
      "1067 steps has passed since last update. Now update the parameter of the behavior policy. Current step: 22127\n",
      "(52.8s,+30.3s)\tIteration 20, current mean episode reward is -19.6. {'loss': 0.5186, 'episode_len': 1066.0}\n",
      "In 20 iteration, current mean episode reward -19.600 is greater than reward threshold -20.0. Congratulation! Now we exit the training process.\n"
     ]
    }
   ],
   "source": [
    "# [optional] BONUS!!! Train DQN in \"Pong-ram-v0\" environment\n",
    "# Tune the hyperparameter and take some time to train agent\n",
    "# You need to install gym[atari] first via `pip install gym[atari]`\n",
    "\n",
    "pytorch_trainer2, _ = run(DQNTrainer, dict(\n",
    "    max_episode_length=10000,\n",
    "    max_iteration=100,\n",
    "    evaluate_interval=10,\n",
    "    evaluate_num_episodes=10,\n",
    "    learning_rate=0.00001,\n",
    "    clip_norm=5.0,\n",
    "    memory_size=1000000,\n",
    "    learn_start=10000,\n",
    "    eps=0.1,\n",
    "    target_update_freq=1000,\n",
    "    learn_freq=4,\n",
    "    batch_size=32,\n",
    "    env_name=\"Pong-ram-v0\"\n",
    "), reward_threshold=-20.0)\n",
    "\n",
    "# This environment is hard to train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [optional] If you have train the agent in Pont-ram-v0, please save the weights so that\n",
    "# we can restore it. Please include the pong-agent.pkl into the zip.\n",
    "\n",
    "import pickle\n",
    "with open(\"pong-agent.pkl\", \"wb\") as f:\n",
    "    pickle.dump(pytorch_trainer2.network.state_dict(), f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average episode reward for your Pytorch agent in Pong-ram-v0:  -19.0\n"
     ]
    }
   ],
   "source": [
    "print(\"Average episode reward for your Pytorch agent in Pong-ram-v0: \",\n",
    "      pytorch_trainer2.evaluate(1, render=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------\n",
    "\n",
    "## Conclusion and Discussion\n",
    "\n",
    "In this assignment, we learn how to build several function approximation algorithm, how to implement basic gradient descent methods and how to use pytorch.\n",
    "\n",
    "It's OK to leave the following cells empty. In the next markdown cell, you can write whatever you like. Like the suggestion on the course, the confusing problems in the assignments, and so on.\n",
    "\n",
    "If you want to do more investigation, feel free to open new cells via `Esc + B` after the next cells and write codes in it, so that you can reuse some result in this notebook. Remember to write sufficient comments and documents to let others know what you are doing.\n",
    "\n",
    "Following the submission instruction in the assignment to submit your assignment to our staff. Thank you!\n",
    "\n",
    "------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "language": "python",
   "name": "torch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
